package com.nlp.trie.viterbi;

import java.awt.Canvas;
import java.awt.Color;
import java.awt.Graphics;
import java.text.DecimalFormat;

public class HMMCanvas extends Canvas{
	/**
	 * 
	 */
	private static final long serialVersionUID = 1L;
	static final int grid_x = 60;
	static final int grid_y = 40;
	static final int offset_x = 70;
	static final int offset_y = 30;
	static final int offset_y2 = 10;
	static final int offset_y3 = 65;
	static final int col_x = 40;
	static final int col_y = 10;
	static final int state_r = 10;
	static final Color state_fill = Color.white;
	static final Color state_fill_maximum = Color.yellow;
	static final Color state_fill_best = Color.red;
	static final Color state_boundery = Color.black;
	static final Color link_normal = Color.green;
	static final Color link_processed = Color.blue;
	static final Color link_maximum = Color.red;
	
	HMMDecoder hmm;
	
	public HMMCanvas() {
		setBackground(Color.white);
		setSize(400, 300);
	}
	
	public void setHmm(HMMDecoder hmm) {
		this.hmm = hmm;
	}
	
	private void drawState(Graphics g, int x, int y, Color c) {
		x = x * grid_x + offset_x;
		y = y * grid_y + offset_y;
		g.setColor(c);
		g.fillOval(x - state_r, y - state_r, state_r * 2, state_r * 2);
		g.setColor(state_boundery);
		g.drawOval(x - state_r, y - state_r, state_r * 2, state_r * 2);
	}
	
	private void drawLink(Graphics g, int x, int y0, int y1, Color c) {
		int x0 = grid_x * x + offset_x;
		int x1 = grid_x * (x + 1) + offset_x;
		y0 = y0 * grid_y + offset_y;
		y1 = y1 * grid_y + offset_y;
		g.setColor(c);
		g.drawLine(x0, y0, x1, y1);
	}
	
	private void drawCenterString(Graphics g, String s, int x, int y) {
		x = x - g.getFontMetrics().stringWidth(s) / 2;
		g.setColor(Color.black);
		g.drawString(s, x, y + 5);
	}

	private void drawRightString(Graphics g, String s, int x, int y) {
		x = x - g.getFontMetrics().stringWidth(s);
		g.setColor(Color.black);
		g.drawString(s, x, y + 5);
	}
	
	public void paint(Graphics g) {
		if (hmm == null) {
			return;
		}
		DecimalFormat form = new DecimalFormat("0.0000");
		int nsymbols = hmm.symbols.size();
		int nstates = hmm.states.size();
		// complete graph.
		for (int i = 0; i < nsymbols; i++) {
			int offset_ymax = offset_y2 + nstates * grid_y;
			if (i < nsymbols - 1) {
				for (int y1 = 0; y1 < nstates; y1++) {
					for (int y0 = 0; y0 < nstates; y0++) {
						Color c = link_normal;
						if (hmm.stage == i + 1 && hmm.i0 == y0 && hmm.i1 == y1) {
							c = link_processed;
						}
						if (hmm.matrix_prevstate[i + 1][y1] == y0) {
							c = link_maximum;
						}
						drawLink(g, i, y0, y1, c);
						if (c == link_maximum && 0 < i) {
							double transprob = hmm.states.get(y0).transprob(
									hmm.states.get(y1));
							drawCenterString(g, form.format(transprob),
									offset_x + i * grid_x + grid_x / 2,
									offset_ymax);
							offset_ymax = offset_ymax + 16;
						}
					}
				}
			}
			// state circles.
			for (int y = 0; y < nstates; y++) {
				Color c = state_fill;
				if (hmm.matrix_prevstate[i][y] != -1) {
					c = state_fill_maximum;
				}
				if (hmm.sequence.size() == nsymbols
						&& hmm.sequence.get(nsymbols - 1 - i) == y) {
					c = state_fill_best;
				}
				drawState(g, i, y, c);
			}
		}
		// max probability.
		for (int i = 0; i < nsymbols; i++) {
			for (int y1 = 0; y1 < nstates; y1++) {
				if (hmm.matrix_prevstate[i][y1] != -1) {
					drawCenterString(g, form.format(hmm.matrix_maxprob[i][y1]),
							offset_x + i * grid_x, offset_y + y1 * grid_y);
				}
			}
		}

		// captions (symbols atop)
		for (int i = 0; i < nsymbols; i++) {
			drawCenterString(g, hmm.symbols.get(i).name, offset_x + i * grid_x,
					col_y);
		}
		// captions (states in left)
		for (int y = 0; y < nstates; y++) {
			drawRightString(g, hmm.states.get(y).name, col_x, offset_y + y
					* grid_y);
		}

		// status bar
		g.setColor(Color.black);
		g.drawString(hmm.status, col_x, offset_y3 + nstates * grid_y -10);
		g.drawString(hmm.status2, col_x, offset_y3 + nstates * grid_y + 16-10);
	}
}
